package com.csx.data.orm.jpa.batch.service;

import cn.hutool.core.collection.CollectionUtil;
import com.csx.data.orm.jpa.base.entity.User;
import com.csx.data.orm.jpa.batch.repositories.UserBatchRepository;
import com.csx.data.orm.jpa.batch.util.SqlUtil;
import javax.persistence.*;
import org.hibernate.Session;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.orm.jpa.EntityManagerFactoryUtils;
import org.springframework.orm.jpa.EntityManagerHolder;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

@Service
public class UserBatchService {

    @Autowired
    private UserBatchRepository userBatchRepository;

    @Autowired
    private EntityManager entityManager;

    @Autowired
    private EntityManagerFactory entityManagerFactory;

    private ThreadLocal<AtomicLong> atomicLongThreadLocal = ThreadLocal.withInitial(() -> new AtomicLong(1));

    @Value("${spring.jpa.batchSize: 200}")
    private int BATCH_SIZE;

    @Autowired
    private JdbcTemplate jdbcTemplate;

    public List<User> saveAll(List<User> list) {
        Iterable<User> users = userBatchRepository.saveAll(list);
        return StreamSupport.stream(users.spliterator(), false)
                .collect(Collectors.toList());
    }

    /**
     * 插入5000行数据需要用 60s 左右
     */
    public List<User> batchExecute(List<User> users) {

        List<User> savedUsers = new ArrayList<>();
        Session session = entityManager.unwrap(Session.class);
        session.doWork(connection -> {
            try (PreparedStatement stmt = connection.prepareStatement("INSERT INTO user(first_name, last_name, age, city) VALUES (?,?,?,?)", Statement.RETURN_GENERATED_KEYS)) {
                int batchSize = 50;
                int count = 0;
                for (User user : users) {
                    stmt.setString(1, user.getFirstName());
                    stmt.setString(2, user.getLastName());
                    stmt.setInt(3, user.getAge());
                    stmt.setString(4, user.getCity());
                    stmt.addBatch();

                    if(++count % batchSize == 0) {
                        stmt.executeBatch();
                        stmt.clearBatch();
                    }

                }
                stmt.executeBatch();

                try (ResultSet rs = stmt.getGeneratedKeys()) {
                    int index = 0;
                    while (rs.next()) {
                        long id = rs.getLong(1);
                        User savedUser = new User(id, users.get(index).getFirstName(),
                                users.get(index).getLastName(),
                                users.get(index).getAge(),
                                users.get(index).getCity());
                        savedUsers.add(savedUser);
                        index++;
                    }
                }
            }
        });
        return savedUsers;
    }

    @Transactional
    public List<User> batchExecuteWithThreadPool(List<User> users) {
        List<User> savedUsers = new ArrayList<>();

        int batchSize = 500; // 每批次的大小
        ExecutorService executorService = Executors.newFixedThreadPool(users.size() / batchSize);

        try {
            for (int i = 0; i < users.size(); i += batchSize) {
                int endIndex = Math.min(i + batchSize, users.size());
                List<User> batchUsers = users.subList(i, endIndex);

                executorService.execute(() -> {
                    EntityManager entityManager = entityManagerFactory.createEntityManager();
                    TransactionSynchronizationManager.bindResource(entityManagerFactory, new EntityManagerHolder(entityManager));
                    EntityTransaction transaction = null;
                    try {

                        transaction = entityManager.getTransaction();
                        transaction.begin();
                        for (User user : batchUsers) {
                            if (user.getId() != null) {
                                entityManager.merge(user);
                            } else {
                                entityManager.persist(user);
                            }
                        }
                        entityManager.flush();
                        entityManager.clear();
                        transaction.commit();
                    } finally {
                        TransactionSynchronizationManager.unbindResource(entityManagerFactory);
                        EntityManagerFactoryUtils.closeEntityManager(entityManager);
                    }
                });
            }

            executorService.shutdown();
            executorService.awaitTermination(1, TimeUnit.MINUTES);

            savedUsers = users;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        return savedUsers;
    }

    @Transactional
    public void batchExecuteWithEntityManager(List<User> users) {

        EntityManagerFactory emf = Persistence.createEntityManagerFactory("my-persistence-unit");
        EntityManager em = emf.createEntityManager();

        final int BATCH_SIZE = 1000;
        ExecutorService executorService = Executors.newFixedThreadPool(users.size() / BATCH_SIZE + 1);
        int saveRetain = users.size();
        if (users.size() > BATCH_SIZE) {
            int start = 0;
            int end = 0;
            while (saveRetain > 0) {

                end = start + Math.min(saveRetain, BATCH_SIZE);
                List<User> subUsersToSave = users.subList(start, end);
                start += BATCH_SIZE;
                executorService.submit(() -> {
                    EntityManager subEm = emf.createEntityManager(); // 创建新的EntityManager
                    try {
                        subEm.getTransaction().begin();
                        List<User> tempUsers = batchSave(subEm, subUsersToSave);
//                    res.addAll(tempUsers);
                        subEm.getTransaction().commit();
                    } catch (Exception e) {

                    } finally {
                        if (subEm != null && subEm.isOpen()) {
                            subEm.close(); // 关闭新的EntityManager
                        }
                    }
                });
                saveRetain -= BATCH_SIZE;
            }
        } else {
            // 直接使用原始的EntityManager
            List<User> tempUsers = batchSave(em, users);
        }

        assert em != null;
        em.close();
        emf.close();


    }


    @Transactional
    public void executeQuery(StringBuilder sb) {
        Session unwrap = entityManager.unwrap(Session.class);
        unwrap.setJdbcBatchSize(1000);
        try {
            unwrap.getTransaction().begin();
            Query query = entityManager.createNativeQuery(sb.toString());
            query.executeUpdate();
            unwrap.getTransaction().commit();
        } catch (Exception e) {
            e.printStackTrace();
            unwrap.getTransaction().rollback();
        }
    }

    @Transactional
    public List<User> batchExecute2(List<User> users) {

        List<User> res = new ArrayList<>();
        final int BATCH_SIZE = 1000;
        ExecutorService executorService = Executors.newFixedThreadPool(users.size() / BATCH_SIZE + 1);
        int saveRetain = users.size();
        if(users.size() > BATCH_SIZE) {
            int start = 0;
            int end = 0;
            while(saveRetain > 0) {

                end = start + Math.min(saveRetain, BATCH_SIZE);
                List<User> subUsersToSave = users.subList(start, end);
                start += BATCH_SIZE;
                executorService.submit(() -> {
                    List<User> tempUsers = batchSave(entityManager, subUsersToSave);
//                    res.addAll(tempUsers);
                });
                saveRetain -= BATCH_SIZE;
            }
        }else {
            List<User> tempUsers = batchSave(entityManager, users);
//            res.addAll(tempUsers);
        }

        return res;
    }

    public List<User> batchSave(EntityManager em, List<User> users) {
        List<User> savedUsers = new ArrayList<>();
        for (User user : users) {
            em.persist(user);
            savedUsers.add(user);
        }
        return savedUsers;
    }

    @Transactional
    public void batchInsertWithJdbcTemplate(List<User> users) {


        if(CollectionUtil.isEmpty(users)) {
            return;
        }
        Iterable<User> all = userBatchRepository.findAllWithLock();
        Optional<Long> maxId = StreamSupport.stream(all.spliterator(), false)
                .map(User::getId)
                .max(Long::compareTo);
        maxId.ifPresent(aLong -> atomicLongThreadLocal.set(new AtomicLong(aLong + 1)));
        String insertSql = SqlUtil.generateInsertSql(users.get(0).getClass());

        // 根据batchSize 分隔 users
        int start = 0;
        int end = 0;
        while(start < users.size()) {

            end = Math.min(start + BATCH_SIZE, users.size());
            List<User> subList = users.subList(start, end);
            start += BATCH_SIZE;

            List<Object[]> args = SqlUtil.generateInsertArgsList(subList, () -> {
                AtomicLong atomicLong = atomicLongThreadLocal.get();
                return atomicLong.getAndIncrement();
            });

            jdbcTemplate.batchUpdate(insertSql, args);
        }
    }

}
